from keras.layers import Conv1D, Dense, BatchNormalization, MaxPooling1D, Activation, Flatten, GlobalAveragePooling1D, \
    Multiply, Lambda, Concatenate
from keras.models import Sequential, Model
from keras.regularizers import l2
import preprocess
from keras.callbacks import TensorBoard
import numpy as np
import tensorflow as tf
from keras import backend as K

# 训练参数
batch_size = 128
epochs = 20
num_classes = 10
length = 2048
BatchNorm = True
number = 1000
normal = True
rate = [0.7, 0.2, 0.1]


# 1. 将动态SE模块改为类形式（解决变量创建问题）
class DynamicMultiscaleSE(tf.keras.layers.Layer):
    def __init__(self, ratio=8, **kwargs):
        super(DynamicMultiscaleSE, self).__init__(**kwargs)
        self.ratio = ratio

    def build(self, input_shape):
        channels = input_shape[-1]
        self.gap = GlobalAveragePooling1D()
        self.dense1 = Dense(channels * 3 // self.ratio, activation='relu')
        self.dense2 = Dense(channels, activation='sigmoid')
        super(DynamicMultiscaleSE, self).build(input_shape)

    def call(self, x):
        # 路径1: 全局平均池化
        gap = self.gap(x)

        # 路径2: 时域峭度
        def get_kurtosis(x):
            mean = tf.reduce_mean(x, axis=1, keepdims=True)
            std = tf.math.reduce_std(x, axis=1, keepdims=True) + 1e-6
            return tf.reduce_mean(((x - mean) / std) ** 4, axis=1) - 3

        kurtosis = Lambda(get_kurtosis)(x)

        # 路径3: 频段能量（简化版，避免复杂FFT）
        def get_energy(x):
            return tf.reduce_mean(tf.abs(x), axis=1)

        energy = Lambda(get_energy)(x)

        # 动态融合
        fused = Concatenate()([gap, kurtosis, energy])
        se = self.dense1(fused)
        se = self.dense2(se)

        return Multiply()([x, se])


# 2. 修改wdcnn函数
def wdcnn(model, filters, kernel_size, strides, conv_padding, pool_padding, pool_size, BatchNormal, use_se=False):
    model.add(Conv1D(filters=filters, kernel_size=kernel_size, strides=strides,
                     padding=conv_padding, kernel_regularizer=l2(1e-4)))
    if BatchNormal:
        model.add(BatchNormalization())
    model.add(Activation('relu'))

    if use_se:
        model.add(DynamicMultiscaleSE(ratio=8))

    model.add(MaxPooling1D(pool_size=pool_size, padding=pool_padding))
    return model


# 数据加载
path = r'C:\Users\xsy\Desktop\wdcnn_bearning_fault_diagnosis-master\data\0HP'
x_train, y_train, x_valid, y_valid, x_test, y_test = preprocess.prepro(
    d_path=path, length=length, number=number, normal=normal, rate=rate, enc=True, enc_step=28)
x_train, x_valid, x_test = x_train[:, :, np.newaxis], x_valid[:, :, np.newaxis], x_test[:, :, np.newaxis]
input_shape = x_train.shape[1:]

# 模型构建
model = Sequential()
model.add(Conv1D(filters=16, kernel_size=64, strides=16, padding='same',
                 kernel_regularizer=l2(1e-4), input_shape=input_shape))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))

# 后续卷积层（在第三层插入SE）
model = wdcnn(model, filters=32, kernel_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm, use_se=False)
model = wdcnn(model, filters=64, kernel_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm, use_se=True)  # 关键修改处
model = wdcnn(model, filters=64, kernel_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm, use_se=False)
model = wdcnn(model, filters=64, kernel_size=3, strides=1, conv_padding='valid',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm, use_se=False)

# 全连接层
model.add(Flatten())
model.add(Dense(units=100, activation='relu', kernel_regularizer=l2(1e-4)))
model.add(Dense(units=num_classes, activation='softmax', kernel_regularizer=l2(1e-4)))

# 编译与训练
model.compile(optimizer='Adam', loss='categorical_crossentropy', metrics=['accuracy'])
tb_cb = TensorBoard(log_dir='logs')
model.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=epochs,
          validation_data=(x_valid, y_valid), callbacks=[tb_cb])

# 评估
score = model.evaluate(x=x_test, y=y_test, verbose=0)
print("测试集损失:", score[0], "测试集准确率:", score[1])

